# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import numpy

from coin.strategy.basis_strat.moving_average import MovingAverage, MovingObject
from coin.experimental.xguo.fruit.base.single_product import SingleProductPriceCalculator


class SmmPricerConfig(object):
  def __init__(self):
    # numbers
    self.window_period = int(300 * 1e9)
    self.enter_position_threshold = 3
    self.book_level = 5

    self.force_fire_threshold = 3 * 1e-4
    self.force_fire_cooldown = 5 * 1e9
    self.force_fire_cooldown_window = 5 * 1e9
    self.price_pull = 2 * 1e-4
    self.price_push = 2 * 1e-4
    self.pos_lean_mul = 0.75
    self.full_pos_lean_mul = 2.
    self.max_lean_size = 10 * 1e-4
    self.force_fire_max_count = 3

    self.force_fire_cap_qty = None
    self.max_pos = None
    self.min_pos = None
    self.min_std_ratio = None
    self.max_std_ratio = None

    # product
    self.trade_symbol = None
    self.ref_symbol = None
    self.get_pos = None


class SmmPricer(object):
  def __init__(self, config):
    self._price_calculators = {}
    self._timestamp = None
    self._config = config
    self._ma = MovingAverage(config.window_period)
    self._agg_buys = MovingObject(config.force_fire_cooldown_window)
    self._agg_sells = MovingObject(config.force_fire_cooldown_window)

    products = (config.trade_symbol, config.ref_symbol)
    for prod in products:
      price_calculator = SingleProductPriceCalculator(prod)
      self._price_calculators[prod] = price_calculator

      if prod == config.trade_symbol:
        self._tpc = price_calculator
      elif prod == config.ref_symbol:
        self._rpc = price_calculator
      else:
        raise ValueError()

  @property
  def trade_symbol(self):
    return self._config.trade_symbol

  @property
  def ref_symbol(self):
    return self._config.ref_symbol

  def get_pos(self):
    return self._config.get_pos()

  def get_norm_position(self):  # normalized to [-1, 1]
    pos_range = self._config.max_pos - self._config.min_pos
    return 1 - 2 * (self.get_pos() - self._config.min_pos) / pos_range

  def update_tick(self):
    self._tpc.renew_tick()
    self._rpc.renew_tick()

  def is_ready(self):
    return self._ma.ready()

  def update(self, product, book):
    self._timestamp = book.timestamp
    self._price_calculators[product].update(book)

    if self._tpc.is_ready() and self._rpc.is_ready():
      # TODO(xguo) Handle no ask/bid/mid case.
      implied_rate = self._tpc.midp / self._rpc.midp
      self._ma.update(self._timestamp, implied_rate)

  def gen_order(self):
    all_orders = {'symbol': str(self.trade_symbol)}
    rate_mean, rate_std = self._ma.get_statistics()
    constrained_std = numpy.clip(
        rate_std,
        self._config.min_std_ratio * rate_mean,
        self._config.max_std_ratio * rate_mean)
    shift_std = constrained_std * self._config.enter_position_threshold
    shift_edge = shift_std / rate_mean
    ref = self._rpc
    trade = self._tpc

    lean_size = self._config.pos_lean_mul * shift_edge * self.get_norm_position()
    capped_lean_size = numpy.clip(lean_size, -self._config.max_lean_size, self._config.max_lean_size)
    if lean_size > 0:
      # Negative position
      raw_pass_buy_price = ref.midp * (rate_mean - shift_std) * (1 + capped_lean_size)
      raw_pass_sell_price = ref.midp * (rate_mean + shift_std) \
          * (1 + lean_size * self._config.full_pos_lean_mul)
    else:
      # Positive position
      raw_pass_sell_price = ref.midp * (rate_mean + shift_std) * (1 + capped_lean_size)
      raw_pass_buy_price = ref.midp * (rate_mean - shift_std) \
          * (1 + lean_size * self._config.full_pos_lean_mul)

    if raw_pass_buy_price * (1 - self._config.force_fire_threshold) > trade.ask1:
      agg_buy_price = trade.get_prev_price(raw_pass_buy_price)
      agg_buy_qty = self._get_force_buy_qty(raw_pass_buy_price, agg_buy_price, trade)
      all_orders['agg_buy'] = dict(price=agg_buy_price, qty=agg_buy_qty)
      self._agg_buys.update(
          self._timestamp,
          (raw_pass_buy_price, agg_buy_price, (trade.ask1, trade.ask1_qty), agg_buy_qty))

    if raw_pass_sell_price * (1 + self._config.force_fire_threshold) < trade.bid1:
      agg_sell_price = trade.get_next_price(raw_pass_sell_price)
      agg_sell_qty = self._get_force_sell_qty(raw_pass_sell_price, agg_sell_price, self._tpc)
      all_orders['agg_sell'] = dict(price=agg_sell_price, qty=agg_sell_qty)
      self._agg_sells.update(
          self._timestamp,
          (raw_pass_sell_price, agg_sell_price, (trade.bid1, trade.bid1_qty), agg_sell_qty))

    pass_buy_price = trade.get_prev_price(min(raw_pass_buy_price, trade.ask1 - 1e-6))
    pass_sell_price = trade.get_next_price(max(raw_pass_sell_price, trade.bid1 + 1e-6))
    buy_price_pull = pass_buy_price * (1 + self._config.price_pull)
    sell_price_pull = pass_sell_price * (1 - self._config.price_pull)
    buy_price_push = pass_buy_price * (1 - self._config.price_push)
    sell_price_push = pass_sell_price * (1 + self._config.price_push)

    all_orders['mm'] = dict(
        pass_sell_price=pass_sell_price,
        pass_buy_price=pass_buy_price,
        sell_price_pull=sell_price_pull,
        buy_price_pull=buy_price_pull,
        sell_price_push=sell_price_push,
        buy_price_push=buy_price_push)
    all_orders['type'] = 'order'
    all_orders['raw_mid_price'] = raw_mid_price
    return all_orders

  def _get_agg_working_qty(self, mo):
    qty = 0.0
    for ts, obj in mo.iteritems():
      qty += obj[3]
    return qty

  def _get_force_buy_qty(self, raw_price, agg_price, pc):
    self._agg_buys.update(self._timestamp)
    size = self._agg_buys.size()
    qty = 0
    working_qty = 0
    avail_qty = 0

    if size >= self._config.force_fire_max_count:
      return 0

    if size <= 0:
      return 0

    last_ts = self._agg_buys.get_latest_ts()
    if self._timestamp - last_ts < self._config.force_fire_cooldown:
      return 0

    last_raw_price, last_agg_price, last_ba, _ = self._agg_buys.get_latest_data()
    # Order book not updated and agg price not improved.
    if (numpy.isclose(last_ba[0], pc.ask0) and
        numpy.isclose(last_ba[1], pc.ask0_qty) and agg_price <= last_agg_price + 1e-6):
      return 0

    # Agg raw price moving lower, do not fire.
    if raw_price < last_raw_price:
      return 0

    working_qty = self._get_agg_working_qty(self._agg_buys)
    for p, q in pc.get_ask_array(self._config.book_level):
      if p > raw_price:
        break
      avail_qty += q
    vacancy_qty = self._config.max_pos - self.get_pos() - working_qty
    qty = min(vacancy_qty, self._config.force_fire_cap_qty, avail_qty)
    return qty

  def _get_force_sell_qty(self, raw_price, agg_price, pc):
    self._agg_sells.update(self._timestamp)
    size = self._agg_sells.size()
    qty = 0
    working_qty = 0
    avail_qty = 0

    if size >= self._config.force_fire_max_count:
      return 0

    if size <= 0:
      return 0

    last_ts = self._agg_sells.get_latest_ts()
    if self._timestamp - last_ts < self._config.force_fire_cooldown:
      return 0

    last_raw_price, last_agg_price, last_bb, _ = self._agg_sells.get_latest_data()
    # Order book not updated and agg price not improved.
    if (numpy.isclose(last_bb[0], pc.bid0) and
        numpy.isclose(last_bb[1], pc.bid0_qty) and agg_price >= last_agg_price - 1e-6):
      return 0

    # Agg raw price moving higher, do not fire.
    if raw_price > last_raw_price:
      return 0

    working_qty = self._get_agg_working_qty(self._agg_sells)
    for b, q in pc.get_bid_array(self._config.book_level):
      if b < raw_price:
        break
      avail_qty += q
    vacancy_qty = self.get_pos() - self._config.min_pos - working_qty
    qty = min(vacancy_qty, self._config.force_fire_cap_qty, avail_qty)
    return qty
